import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

# 1. 数据预处理（转换为Tensor并归一化）
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  # 归一化到 [-1, 1]

# 2. 加载CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64, shuffle=False, num_workers=2)

class_names = ('plane', 'car', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck')

# 3. 可视化部分训练数据
def imshow(img):
    img = img / 2 + 0.5  # 反归一化（从[-1,1]到[0,1]）
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 转换通道顺序
    plt.show()

# 随机获取一些训练图片
dataiter = iter(trainloader)
images, labels = next(dataiter)

# 显示图片
imshow(torchvision.utils.make_grid(images[:5]))
# 打印标签
print(' '.join(f'{class_names[labels[j]]:5s}' for j in range(5)))

# 4. 定义CNN模型
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)  # 输入3通道，32个3x3卷积核
        self.pool = nn.MaxPool2d(2, 2)  # 2x2池化
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.fc1 = nn.Linear(64 * 4 * 4, 64)  # 展平后连接全连接层
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = F.relu(self.conv3(x))
        x = torch.flatten(x, 1)  # 展平（保留batch维度）
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = Net()

# 5. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失（适用于分类）
optimizer = optim.Adam(net.parameters(), lr=0.001)  # Adam优化器

# 6. 训练模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)  # 移动模型到GPU（如果可用）

train_losses = []
train_accs = []
val_losses = []
val_accs = []

for epoch in range(10):  # 训练10轮
    running_loss = 0.0
    correct = 0
    total = 0
    net.train()  # 训练模式
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 统计训练数据
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    # 计算训练集准确率和损失
    train_acc = correct / total
    train_loss = running_loss / len(trainloader)
    train_accs.append(train_acc)
    train_losses.append(train_loss)
    
    # 在验证集上评估
    net.eval()  # 评估模式
    val_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():  # 关闭梯度计算
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    val_acc = correct / total
    val_loss /= len(testloader)
    val_accs.append(val_acc)
    val_losses.append(val_loss)
    
    print(f'Epoch {epoch+1}, 训练损失: {train_loss:.4f}, 训练准确率: {train_acc:.4f}, 验证损失: {val_loss:.4f}, 验证准确率: {val_acc:.4f}')

print('训练完成')

# 7. 可视化训练过程
plt.figure(figsize=(12,4))
# 准确率曲线
plt.subplot(1,2,1)
plt.plot(train_accs, label='训练准确率')
plt.plot(val_accs, label='验证准确率')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

# 损失曲线
plt.subplot(1,2,2)
plt.plot(train_losses, label='训练损失')
plt.plot(val_losses, label='验证损失')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 8. 保存模型
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)